import logging
import numpy as np

from util.util_log import setup_logging
from FitnessMultiBeamTrace import FitnessMultiBeamTrace


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[BDE]")



class BinaryDifferentialEvolution:
    def __init__(self, objective_function, dimensions, population_size=50, F=0.5, CR=0.7, generations=100):
        """
        初始化二进制差分进化算法。

        :param objective_function: 适应度函数
        :param dimensions: 搜索空间维度 (tuple of two integers for 2D arrays)
        :param population_size: 种群大小
        :param F: 差分权重因子
        :param CR: 交叉概率
        :param generations: 迭代次数
        """
        self.objective_function = objective_function
        self.dimensions = dimensions
        self.population_size = population_size
        self.F = F
        self.CR = CR
        self.generations = generations
        self.population = []
        self.best_individual = None
        self.best_fitness = float('-inf')

    def __generate_population_part(self, before_array, current_array, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(before_array.copy().astype(int))
        part.append(current_array.copy().astype(int))
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array 并强制转换为 int 类型
            modified_array = np.copy(current_array).astype(int)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = before_array[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(modified_array)
        return part

    def initialize_population(self):
        """初始化种群"""
        self.population = [np.random.randint(2, size=self.dimensions) for _ in range(self.population_size)]

    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before,
                                    phaseBit_mix_qe):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps, self.population_size // 2)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                 self.population_size - self.population_size // 2)
        self.population = parts_ps + parts_qe

    def differential_mutation(self, target_index):
        """差分变异操作"""
        indices = list(range(self.population_size))
        indices.remove(target_index)
        a, b, c = np.random.choice(indices, 3, replace=False)
        random_mask = (np.random.rand(*self.dimensions) < self.F).astype(int)
        mutant = self.population[a] ^ ((self.population[b] ^ self.population[c]) & random_mask)
        return mutant

    def crossover(self, target, mutant):
        """交叉操作：二项式交叉"""
        trial = np.copy(target)
        j_rand = np.random.randint(np.prod(self.dimensions))
        for j in range(np.prod(self.dimensions)):
            if np.random.rand() < self.CR or j == j_rand:
                trial[j // self.dimensions[1], j % self.dimensions[1]] = mutant[
                    j // self.dimensions[1], j % self.dimensions[1]]
        return trial

    def evaluate_population(self):
        """评估种群中每个个体的适应度值"""
        fitness_values = np.array([self.objective_function(individual) for individual in self.population])
        return fitness_values

    def run(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe):
        """运行二进制差分进化算法"""

        # """初始化返回值"""
        self.best_individual = phaseBit_mix_ps
        self.best_fitness = self.objective_function(self.best_individual)
        self.best_individual_history = []
        self.best_fitness_history = []

        # 普通初始化
        # self.initialize_population()

        # 初始化
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe)

        for generation in range(self.generations):
            new_population = []

            for i in range(self.population_size):
                target = self.population[i]
                mutant = self.differential_mutation(i)
                trial = self.crossover(target, mutant)

                # Evaluate the trial vector
                trial_fitness = self.objective_function(trial)
                target_fitness = self.objective_function(target)

                # Selection: Replace target with trial if trial is better
                if trial_fitness > target_fitness:
                    new_population.append(trial)
                else:
                    new_population.append(target)

            # Update the population
            self.population = new_population

            # Find the best individual and its fitness
            fitness_values = self.evaluate_population()
            best_index = np.argmax(fitness_values)
            current_best_individual = self.population[best_index]
            current_best_fitness = fitness_values[best_index]

            # Update global best if necessary
            if current_best_fitness > self.best_fitness:
                self.best_fitness = current_best_fitness
                self.best_individual = current_best_individual

            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # Print current best fitness
            # print(f"Generation {generation + 1}: Best Fitness = {self.best_fitness}")
            logger.info("iteration=%d: self.best_fitness=%f" % (generation, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")

        # return self.best_individual, self.best_fitness
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history


# 示例目标函数：最大化一的数量
def maximize_ones(matrix):
    return np.sum(matrix)


def test():
    # 使用二进制差分进化算法优化最大化一的数量的问题
    dimensions = (4, 4)  # Example 3x3 matrix
    population_size = 50
    F = 0.5
    CR = 0.7
    generations = 100
    # 创建一个示例 FitnessMultiBeamTrace 实例
    bit_num = 8  # 根据实际情况设置比特数
    beam_num = 2  # 根据实际情况设置波束数
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)

    # 初始化主瓣位置
    phaseBit_target = np.random.randint(2, size=(4, 4))
    fitness_function.init_main_lope_pos(phaseBit_target)

    # 定义初始相位比特矩阵
    phaseBit_ps_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_ps = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1]])
    phaseBit_qe_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_qe = np.array([[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])

    # 使用自定义的适应度函数
    bde_custom_fitness = BinaryDifferentialEvolution(objective_function=fitness_function.fitness,
                                                     dimensions=dimensions,
                                                     population_size=population_size,
                                                     F=F,
                                                     CR=CR,
                                                     generations=generations)
    best_individual_custom, best_fitness_custom, best_fitness_history_custom, best_individual_history_custom \
        = bde_custom_fitness.run(phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)

    print(f"\nFinal Best Individual with Custom Fitness Function:\n{best_individual_custom}")
    print(f"Final Best Fitness with Custom Fitness Function: {best_fitness_custom}")


if __name__=="__main__":
    test()